"""
音频格式转换和优化模块
"""

import io
import tempfile
from typing import Optional, Tuple, Dict, Any
import numpy as np
import soundfile as sf
from pydub import AudioSegment
from scipy import signal
import librosa


class AudioConverter:
    """音频格式转换器"""
    
    def __init__(self):
        self.supported_formats = {
            'wav', 'mp3', 'ogg', 'flac', 'aac', 'm4a'
        }
    
    def convert_format(
        self, 
        input_data: bytes, 
        input_format: str,
        output_format: str,
        sample_rate: Optional[int] = None,
        bitrate: str = "192k"
    ) -> bytes:
        """转换音频格式"""
        if input_format not in self.supported_formats:
            raise ValueError(f"Unsupported input format: {input_format}")
        
        if output_format not in self.supported_formats:
            raise ValueError(f"Unsupported output format: {output_format}")
        
        try:
            # 使用pydub进行格式转换
            if input_format == 'wav':
                audio = AudioSegment.from_wav(io.BytesIO(input_data))
            elif input_format == 'mp3':
                audio = AudioSegment.from_mp3(io.BytesIO(input_data))
            elif input_format == 'ogg':
                audio = AudioSegment.from_ogg(io.BytesIO(input_data))
            elif input_format == 'flac':
                audio = AudioSegment.from_flac(io.BytesIO(input_data))
            else:
                # 对于其他格式，先用soundfile读取
                with tempfile.NamedTemporaryFile() as temp_file:
                    temp_file.write(input_data)
                    temp_file.seek(0)
                    
                    data, sr = sf.read(temp_file.name)
                    audio = AudioSegment(
                        data.tobytes(),
                        frame_rate=sr,
                        sample_width=data.dtype.itemsize,
                        channels=1 if len(data.shape) == 1 else data.shape[1]
                    )
            
            # 转换采样率
            if sample_rate:
                audio = audio.set_frame_rate(sample_rate)
            
            # 导出为目标格式
            output_buffer = io.BytesIO()
            
            if output_format == 'mp3':
                audio.export(output_buffer, format='mp3', bitrate=bitrate)
            elif output_format == 'wav':
                audio.export(output_buffer, format='wav')
            elif output_format == 'ogg':
                audio.export(output_buffer, format='ogg')
            elif output_format == 'flac':
                audio.export(output_buffer, format='flac')
            else:
                audio.export(output_buffer, format=output_format)
            
            output_buffer.seek(0)
            return output_buffer.read()
            
        except Exception as e:
            raise RuntimeError(f"Audio conversion failed: {str(e)}")
    
    def optimize_audio(
        self,
        audio_data: bytes,
        format: str = 'wav',
        target_sample_rate: int = 22050,
        normalize: bool = True,
        trim_silence: bool = True
    ) -> bytes:
        """优化音频质量"""
        try:
            # 读取音频
            with tempfile.NamedTemporaryFile() as temp_file:
                temp_file.write(audio_data)
                temp_file.seek(0)
                
                # 使用librosa进行高级处理
                y, sr = librosa.load(temp_file.name, sr=None)
            
            # 转换采样率
            if target_sample_rate != sr:
                y = librosa.resample(y, orig_sr=sr, target_sr=target_sample_rate)
                sr = target_sample_rate
            
            # 标准化音量
            if normalize:
                y = librosa.util.normalize(y)
            
            # 去除静音
            if trim_silence:
                y, _ = librosa.effects.trim(y, top_db=20)
            
            # 保存处理后的音频
            with tempfile.NamedTemporaryFile() as temp_file:
                sf.write(temp_file.name, y, sr, format=format)
                with open(temp_file.name, 'rb') as f:
                    return f.read()
                    
        except Exception as e:
            raise RuntimeError(f"Audio optimization failed: {str(e)}")


class AudioProcessor:
    """音频处理器"""
    
    def __init__(self):
        self.converter = AudioConverter()
    
    def apply_effects(
        self,
        audio_data: bytes,
        effects: Dict[str, Any]
    ) -> bytes:
        """应用音频效果"""
        try:
            # 加载音频
            with tempfile.NamedTemporaryFile() as temp_file:
                temp_file.write(audio_data)
                temp_file.seek(0)
                
                y, sr = sf.read(temp_file.name)
            
            # 应用各种效果
            if effects.get('pitch_shift'):
                n_steps = effects['pitch_shift']
                y = librosa.effects.pitch_shift(y, sr=sr, n_steps=n_steps)
            
            if effects.get('time_stretch'):
                rate = effects['time_stretch']
                y = librosa.effects.time_stretch(y, rate=rate)
            
            if effects.get('noise_reduction'):
                y = self._reduce_noise(y, sr)
            
            if effects.get('equalizer'):
                y = self._apply_equalizer(y, sr, effects['equalizer'])
            
            # 保存结果
            with tempfile.NamedTemporaryFile() as temp_file:
                sf.write(temp_file.name, y, sr)
                with open(temp_file.name, 'rb') as f:
                    return f.read()
                    
        except Exception as e:
            raise RuntimeError(f"Audio effect application failed: {str(e)}")
    
    def _reduce_noise(self, y: np.ndarray, sr: int) -> np.ndarray:
        """简单的噪声抑制"""
        # 使用谱减法进行简单降噪
        stft = librosa.stft(y)
        magnitude = np.abs(stft)
        phase = np.angle(stft)
        
        # 估计噪声谱（使用前几帧）
        noise_magnitude = np.mean(magnitude[:, :10], axis=1, keepdims=True)
        
        # 谱减法
        alpha = 2.0  # 过减因子
        magnitude_denoised = magnitude - alpha * noise_magnitude
        magnitude_denoised = np.maximum(magnitude_denoised, 0.1 * magnitude)
        
        # 重构信号
        stft_denoised = magnitude_denoised * np.exp(1j * phase)
        y_denoised = librosa.istft(stft_denoised)
        
        return y_denoised
    
    def _apply_equalizer(
        self, 
        y: np.ndarray, 
        sr: int, 
        eq_settings: Dict[str, float]
    ) -> np.ndarray:
        """应用均衡器"""
        # 简单的三段均衡器
        low_gain = eq_settings.get('low', 1.0)
        mid_gain = eq_settings.get('mid', 1.0)
        high_gain = eq_settings.get('high', 1.0)
        
        # 设计滤波器
        nyquist = sr // 2
        low_cutoff = 300 / nyquist
        high_cutoff = 3000 / nyquist
        
        # 低频滤波器
        b_low, a_low = signal.butter(2, low_cutoff, btype='low')
        low_freq = signal.filtfilt(b_low, a_low, y)
        
        # 高频滤波器
        b_high, a_high = signal.butter(2, high_cutoff, btype='high')
        high_freq = signal.filtfilt(b_high, a_high, y)
        
        # 中频 = 原信号 - 低频 - 高频
        mid_freq = y - low_freq - high_freq
        
        # 应用增益并混合
        y_eq = (low_freq * low_gain + 
                mid_freq * mid_gain + 
                high_freq * high_gain)
        
        return y_eq
    
    def get_audio_info(self, audio_data: bytes) -> Dict[str, Any]:
        """获取音频信息"""
        try:
            with tempfile.NamedTemporaryFile() as temp_file:
                temp_file.write(audio_data)
                temp_file.seek(0)
                
                y, sr = sf.read(temp_file.name)
                
                duration = len(y) / sr
                rms = np.sqrt(np.mean(y**2))
                peak = np.max(np.abs(y))
                
                return {
                    'duration': duration,
                    'sample_rate': sr,
                    'channels': 1 if len(y.shape) == 1 else y.shape[1],
                    'rms_level': float(rms),
                    'peak_level': float(peak),
                    'samples': len(y)
                }
        except Exception as e:
            raise RuntimeError(f"Failed to get audio info: {str(e)}")
